import cv2

# 加载训练数据集文件
recognizer = cv2.face.LBPHFaceRecognizer_create()
recognizer.read('trainer/trainer.yml')

# 创建人脸检测分类器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# 打开摄像头
video_capture = cv2.VideoCapture(0)

# 设置字体参数
font_path = "ZiTiQuanWeiJunHei-W2-2.ttf"  # 汉字字体文件路径
font_size = 0.9
font_color = (0, 255, 0)  # 绿色
thickness = 2

while True:
    # 读取视频帧
    ret, frame = video_capture.read()

    # 将视频帧转换为灰度图像
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 进行人脸检测
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    # 遍历检测到的人脸
    for (x, y, w, h) in faces:
        # 绘制人脸框
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

        # 人脸识别
        roi_gray = gray[y:y + h, x:x + w]
        id, confidence = recognizer.predict(roi_gray)
        if confidence < 80:
            # 显示标签和置信评分
            if id <=7:
                label = "Person  YZH " + str(id)
            elif id >=8 and id <=15:
                label = "Person  RXY " + str(id)
            elif id >=16 and id <=23:
                label = "Person WD" + str(id)
            else:
                label = "Person XXX" + str(id)
            confidence_text = "{0}%".format(round(100 - confidence))
            cv2.putText(frame, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, font_size, font_color, thickness,
                        cv2.LINE_AA)
            cv2.putText(frame, confidence_text, (x, y + h + 25), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    # 显示视频帧
    cv2.imshow('Video', frame)

    # 等待退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头和窗口
video_capture.release()
cv2.destroyAllWindows()